{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CelebA.ipynb",
      "provenance": [
        {
          "file_id": "1pa1835OavB8rJ2HP2FdC6AkoVymzmodm",
          "timestamp": 1622581117510
        },
        {
          "file_id": "1sXJeXxz1r2hW_qtnY6J0aFM5qxWDKj39",
          "timestamp": 1611905139543
        },
        {
          "file_id": "11xlt3cH_t8sdQJIkYc1wN-A8xtzBqj_g",
          "timestamp": 1600557824757
        },
        {
          "file_id": "1OplIUXiVWHjmeXr8A0FMe4kH8NLN1JcD",
          "timestamp": 1598651654838
        },
        {
          "file_id": "1mvp25SrzH1w4f8qrAnTHFZXJHZkL_JuN",
          "timestamp": 1589231089092
        }
      ],
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      }
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SGuSrWEYYUjr"
      },
      "source": [
        "### CelebA Experiments for...\n",
        "##Efficiently Identifying Task Groupings for Multi-Task Learning\n",
        "\n",
        "Licensed under the Apache License, Version 2.0\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r14fs3V2XLOl"
      },
      "source": [
        "import itertools\n",
        "import pickle\n",
        "import time\n",
        "import copy\n",
        "import math\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "from collections import namedtuple, OrderedDict\n",
        "from tqdm import tqdm\n",
        "\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras.initializers import glorot_uniform\n",
        "from tensorflow.keras.layers import Activation, Add, AveragePooling2D, BatchNormalization, Conv2D, Dense, Flatten, MaxPooling2D, Lambda\n",
        "\n",
        "import scipy.integrate as it\n",
        "from absl import flags"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gjx4YgM6giJJ"
      },
      "source": [
        "# Adapted from https://github.com/tianheyu927/PCGrad/blob/master/PCGrad_tf.py\n",
        "GATE_OP = 1\n",
        "\n",
        "class PCGrad(tf.compat.v1.train.Optimizer):\n",
        "  \"\"\"PCGrad. https://arxiv.org/pdf/2001.06782.pdf.\"\"\"\n",
        "\n",
        "  def __init__(self, opt, use_locking=False, name=\"PCGrad\"):\n",
        "    \"\"\"optimizer: the optimizer being wrapped.\"\"\"\n",
        "    super(PCGrad, self).__init__(use_locking, name)\n",
        "    self.optimizer = opt\n",
        "\n",
        "  def compute_gradients(self, loss, var_list=None,\n",
        "                        gate_gradients=GATE_OP,\n",
        "                        aggregation_method=None,\n",
        "                        colocate_gradients_with_ops=False,\n",
        "                        grad_loss=None):\n",
        "    assert isinstance(loss, list)\n",
        "    num_tasks = len(loss)\n",
        "    loss = tf.stack(loss)\n",
        "    tf.random.shuffle(loss)\n",
        "\n",
        "    # Compute per-task gradients.\n",
        "    grads_task = tf.vectorized_map(lambda x: tf.concat(\n",
        "        [tf.reshape(grad, [-1,]) for grad in tf.gradients(\n",
        "            x, var_list) if grad is not None], axis=0), loss)\n",
        "\n",
        "    # Compute gradient projections.\n",
        "    def proj_grad(grad_task):\n",
        "      for k in range(num_tasks):\n",
        "        inner_product = tf.reduce_sum(grad_task*grads_task[k])\n",
        "        proj_direction = inner_product / tf.reduce_sum(\n",
        "            grads_task[k]*grads_task[k])\n",
        "        grad_task = grad_task - tf.minimum(proj_direction, 0.) * grads_task[k]\n",
        "      return grad_task\n",
        "\n",
        "    proj_grads_flatten = tf.vectorized_map(proj_grad, grads_task)\n",
        "\n",
        "    # Unpack flattened projected gradients back to their original shapes.\n",
        "    proj_grads = []\n",
        "    for j in range(num_tasks):\n",
        "      start_idx = 0\n",
        "      for idx, var in enumerate(var_list):\n",
        "        grad_shape = var.get_shape()\n",
        "        flatten_dim = np.prod(\n",
        "            [grad_shape.dims[i].value for i in range(len(grad_shape.dims))])\n",
        "        proj_grad = proj_grads_flatten[j][start_idx:start_idx+flatten_dim]\n",
        "        proj_grad = tf.reshape(proj_grad, grad_shape)\n",
        "        if len(proj_grads) < len(var_list):\n",
        "          proj_grads.append(proj_grad)\n",
        "        else:\n",
        "          proj_grads[idx] += proj_grad\n",
        "        start_idx += flatten_dim\n",
        "    grads_and_vars = list(zip(proj_grads, var_list))\n",
        "    return grads_and_vars\n",
        "\n",
        "  def _create_slots(self, var_list):\n",
        "    self.optimizer._create_slots(var_list)\n",
        "\n",
        "  def _prepare(self):\n",
        "    self.optimizer._prepare()\n",
        "\n",
        "  def _apply_dense(self, grad, var):\n",
        "    return self.optimizer._apply_dense(grad, var)\n",
        "\n",
        "  def _resource_apply_dense(self, grad, var):\n",
        "    return self.optimizer._resource_apply_dense(grad, var)\n",
        "\n",
        "  def _apply_sparse_shared(self, grad, var, indices, scatter_add):\n",
        "    return self.optimizer._apply_sparse_shared(grad, var, indices, scatter_add)\n",
        "\n",
        "  def _apply_sparse(self, grad, var):\n",
        "    return self.optimizer._apply_sparse(grad, var)\n",
        "\n",
        "  def _resource_scatter_add(self, x, i, v):\n",
        "    return self.optimizer._resource_scatter_add(x, i, v)\n",
        "\n",
        "  def _resource_apply_sparse(self, grad, var, indices):\n",
        "    return self.optimizer._resource_apply_sparse(grad, var, indices)\n",
        "\n",
        "  def _finish(self, update_ops, name_scope):\n",
        "    return self.optimizer._finish(update_ops, name_scope)\n",
        "\n",
        "  def _call_if_callable(self, param):\n",
        "    \"\"\"Call the function if param is callable.\"\"\"\n",
        "    return param() if callable(param) else param"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ooXuDeQVe_MH"
      },
      "source": [
        "def del_all_flags(FLAGS):\n",
        "    flags_dict = FLAGS._flags()    \n",
        "    keys_list = [keys for keys in flags_dict]    \n",
        "    for keys in keys_list:\n",
        "        FLAGS.__delattr__(keys)\n",
        "del_all_flags(flags.FLAGS)\n",
        "\n",
        "FLAGS = flags.FLAGS\n",
        "\n",
        "flags.DEFINE_integer('steps', 100, 'Number of epoch to train.')\n",
        "flags.DEFINE_integer('batch_size', 256, 'Number of examples in a minibatch.')\n",
        "flags.DEFINE_integer('order', -1, 'Order of permutations to consider.')\n",
        "flags.DEFINE_enum('eval', 'test', ['valid', 'test'], 'The eval dataset.')\n",
        "flags.DEFINE_enum('method', 'mtl', ['mtl', 'fast_mtl'],'Multitask Training Method.')\n",
        "flags.DEFINE_list('tasks', ['5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', \n",
        "                   'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose',\n",
        "                   'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', \n",
        "                   'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair', \n",
        "                   'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open',\n",
        "                   'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin',\n",
        "                   'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns', \n",
        "                   'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings', \n",
        "                   'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', \n",
        "                   'Wearing_Necktie', 'Young'], \"The attributes to predict in CelebA.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jTH-4Vc5OX-F"
      },
      "source": [
        "SEED = 0\n",
        "METRICS_AVERAGE = 1\n",
        "EPSILON = 0.001\n",
        "TRAIN_SIZE = 162770\n",
        "VALID_SIZE = 19867\n",
        "TEST_SIZE = 19962"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nArMmed5iZmC"
      },
      "source": [
        "class ResBlock(tf.keras.Model):\n",
        "\n",
        "  def __init__(self, filters, kernel_size, strides, name):\n",
        "    super(ResBlock, self).__init__()\n",
        "    self.conv1 = Conv2D(\n",
        "        filters=filters[0],\n",
        "        kernel_size=kernel_size[0],\n",
        "        strides=strides,\n",
        "        name='conv{}_1'.format(name),\n",
        "        kernel_initializer=glorot_uniform(seed=SEED),\n",
        "        padding='same',\n",
        "        use_bias=False)\n",
        "    self.bn1 = BatchNormalization(axis=3, name='bn{}_1'.format(name))\n",
        "    self.conv2 = Conv2D(\n",
        "        filters=filters[1],\n",
        "        kernel_size=kernel_size[1],\n",
        "        strides=(1,1),\n",
        "        name='conv{}_2'.format(name),\n",
        "        kernel_initializer=glorot_uniform(seed=SEED),\n",
        "        padding='same',\n",
        "        use_bias=False)\n",
        "    self.bn2 = BatchNormalization(axis=3, name='bn{}_2'.format(name))\n",
        "\n",
        "    if strides == (1,1):\n",
        "      self.shortcut = Lambda(lambda x : x)\n",
        "    else:\n",
        "      self.shortcut = tf.keras.Sequential()\n",
        "      shortcut_conv = Conv2D(filters=filters[1], \n",
        "                             kernel_size=1, \n",
        "                             strides=(2,2), \n",
        "                             name='skip_conv{}_1'.format(name), \n",
        "                             kernel_initializer=glorot_uniform(seed=SEED),\n",
        "                             padding='valid',\n",
        "                             use_bias=False)\n",
        "      shortcut_bn = BatchNormalization(axis=3, name='skip_bn{}_1'.format(name))\n",
        "      self.shortcut.add(shortcut_conv)\n",
        "      self.shortcut.add(shortcut_bn)\n",
        "\n",
        "  def call(self, inputs):\n",
        "    x = inputs\n",
        "    x = Activation('relu')(self.bn1(self.conv1(x)))\n",
        "    x = self.bn2(self.conv2(x))\n",
        "    x = Add()([x, self.shortcut(inputs)])\n",
        "    return Activation('relu')(x)\n",
        "\n",
        "class ResNet18(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(ResNet18, self).__init__()\n",
        "    self.conv1_1 = Conv2D(\n",
        "        filters=64,\n",
        "        kernel_size=3,\n",
        "        strides=(1, 1),\n",
        "        name='conv1_1',\n",
        "        kernel_initializer=glorot_uniform(seed=SEED),\n",
        "        padding='same',\n",
        "        use_bias=False)\n",
        "    self.bn1_1 = BatchNormalization(axis=3, name='bn1_1')\n",
        "    self.resblock_2 = ResBlock([64, 64], [3, 3], (1, 1), '1')\n",
        "\n",
        "  def call(self, inputs):\n",
        "    x = inputs\n",
        "    x = Activation('relu')(self.bn1_1(self.conv1_1(x)))\n",
        "    x = MaxPooling2D((3,3), strides=(2,2))(x)\n",
        "    x = self.resblock_2(x)\n",
        "    x = AveragePooling2D((2,2), name='avg_pool')(x)\n",
        "    x = Flatten()(x)\n",
        "    return x\n",
        "\n",
        "class AttributeDecoder(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(AttributeDecoder, self).__init__()\n",
        "    self.fc1 = Dense(2, kernel_initializer=glorot_uniform(seed=SEED))\n",
        "\n",
        "  def call(self, inputs):\n",
        "    x = inputs\n",
        "    x = self.fc1(x)\n",
        "    return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IvFDpI8pDdxq"
      },
      "source": [
        "def res_block_step(inputs, base_updated):\n",
        "  conv1 = tf.nn.conv2d(inputs, base_updated[0], strides=(2,2), padding=\"SAME\")\n",
        "  mean1, variance1 = tf.nn.moments(conv1, axes=[0,1,2])\n",
        "  gamma1, beta1 = base_updated[1], base_updated[2]\n",
        "  bn_conv1 = tf.nn.batch_normalization(conv1, mean1, variance1, offset=beta1, scale=gamma1, variance_epsilon=EPSILON)\n",
        "  relu1 = tf.nn.relu(bn_conv1)\n",
        "  \n",
        "  conv2 = tf.nn.conv2d(relu1, base_updated[3], strides=(1,1), padding=\"SAME\")\n",
        "  mean2, variance2 = tf.nn.moments(conv2, axes=[0,1,2])\n",
        "  gamma2, beta2 = base_updated[4], base_updated[5]\n",
        "  bn_conv2 = tf.nn.batch_normalization(conv2, mean2, variance2, offset=beta2, scale=gamma2, variance_epsilon=EPSILON)\n",
        "\n",
        "  skip_conv = tf.nn.conv2d(inputs, base_updated[6], strides=(2,2), padding=\"VALID\")\n",
        "  skip_mean, skip_variance = tf.nn.moments(skip_conv, axes=[0,1,2])\n",
        "  skip_gamma, skip_beta = base_updated[7], base_updated[8]\n",
        "  skip_bn = tf.nn.batch_normalization(skip_conv, skip_mean, skip_variance, offset=skip_beta, scale=skip_gamma, variance_epsilon=EPSILON)\n",
        "  \n",
        "  res_block = tf.nn.relu(bn_conv2 + skip_bn)\n",
        "  return res_block\n",
        "\n",
        "def base_step(inputs, base_updated):\n",
        "  # ResNet Block 1 Output.\n",
        "  conv1_1 = tf.nn.conv2d(inputs, base_updated[0], strides=(1,1), padding=\"SAME\")\n",
        "  mean1_1, variance1_1 = tf.nn.moments(conv1_1, axes=[0,1,2], keepdims=True) # normalize across the channel dimension for spacial batch norm.\n",
        "  gamma1_1, beta1_1 = base_updated[1], base_updated[2]\n",
        "  bn_conv1_1 = tf.nn.batch_normalization(conv1_1, mean1_1, variance1_1, offset=beta1_1, scale=gamma1_1, variance_epsilon=EPSILON)\n",
        "  res_block_1 = tf.nn.max_pool2d(tf.nn.relu(bn_conv1_1), ksize=[1,3,3,1], strides=[1,2,2,1], padding=\"VALID\")\n",
        "\n",
        "  # ResNet Block 2\n",
        "  conv2_1 = tf.nn.conv2d(res_block_1, base_updated[3], strides=(1,1), padding=\"SAME\")\n",
        "  mean2_1, variance2_1 = tf.nn.moments(conv2_1, axes=[0,1,2])\n",
        "  gamma2_1, beta2_1 = base_updated[4], base_updated[5]\n",
        "  bn_conv2_1 = tf.nn.batch_normalization(conv2_1, mean2_1, variance2_1, offset=beta2_1, scale=gamma2_1, variance_epsilon=EPSILON)\n",
        "  res_block2_1 = tf.nn.relu(bn_conv2_1)\n",
        "  \n",
        "  conv2_2 = tf.nn.conv2d(res_block2_1, base_updated[6], strides=(1,1), padding=\"SAME\")\n",
        "  mean2_2, variance2_2 = tf.nn.moments(conv2_2, axes=[0,1,2])\n",
        "  gamma2_2, beta2_2 = base_updated[7], base_updated[8]\n",
        "  bn_conv2_2 = tf.nn.batch_normalization(conv2_2, mean2_2, variance2_2, offset=beta2_2, scale=gamma2_2, variance_epsilon=EPSILON)\n",
        "  res_block_2 = tf.nn.relu(bn_conv2_2 + res_block_1)\n",
        "  \n",
        "  avg_pool = tf.nn.avg_pool2d(res_block_2, ksize=[1,2,2,1], strides=[1,2,2,1], padding=\"VALID\")\n",
        "  resnet_out = tf.reshape(avg_pool, [inputs.shape[0], -1])\n",
        "  return resnet_out"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5yJVRK0cWaYK"
      },
      "source": [
        "def permute(losses):\n",
        "  \"\"\"Returns all combinations of losses in the loss dictionary.\"\"\"\n",
        "  losses = OrderedDict(sorted(losses.items()))\n",
        "  rtn = {}\n",
        "  for task,loss in losses.items():\n",
        "    tmp_dict = {task:loss}\n",
        "    for saved_task, saved_loss in rtn.items():\n",
        "      if FLAGS.order == 1:\n",
        "        continue # Skip higher than first-order combinations.\n",
        "      new_task = \"{}_{}\".format(saved_task, task)\n",
        "      new_loss = loss + saved_loss \n",
        "      tmp_dict[new_task] = new_loss \n",
        "    rtn.update(tmp_dict)\n",
        "  \n",
        "  if FLAGS.order == 1:\n",
        "    rtn[\"_\".join(losses.keys())] = sum(losses.values())\n",
        "  return rtn\n",
        "\n",
        "\n",
        "def permute_list(lst):\n",
        "  \"\"\"Returns all combinations of tasks in the task list.\"\"\"\n",
        "  lst.sort()\n",
        "  rtn = []\n",
        "  for task in lst:\n",
        "    tmp_lst = [task]\n",
        "    for saved_task in rtn:\n",
        "      if FLAGS.order == 1:\n",
        "        continue\n",
        "      new_task = \"{}_{}\".format(saved_task, task)\n",
        "      tmp_lst.append(new_task)\n",
        "    rtn += tmp_lst\n",
        "  \n",
        "  if FLAGS.order == 1:\n",
        "    rtn.append(\"_\".join(lst))\n",
        "  return rtn\n",
        "\n",
        "def decay_lr(step, optimizer):\n",
        "  if (step + 1) % 15 == 0:\n",
        "    optimizer.lr = optimizer.lr / 2.\n",
        "    print('Decreasing the learning rate by 1/2. New Learning Rate: {}'.format(optimizer.lr))\n",
        "\n",
        "def decay_pcgrad_lr(step, lr_var):\n",
        "  if (step + 1) % 15 == 0:\n",
        "    lr_var.assign(lr_var / 2.)\n",
        "    print('Decreasing the learning rate by 1/2.')\n",
        "\n",
        "def add_average(lst, metrics_dict, n):\n",
        "  if len(lst) < n:\n",
        "    lst.append(metrics_dict)\n",
        "  elif len(lst) == n:\n",
        "    lst.pop(0)\n",
        "    lst.append(metrics_dict)\n",
        "  elif len(lst) > n:\n",
        "    raise Exception('List size is greater than n. This should never happen.')\n",
        "\n",
        "def compute_average(metrics_list, n):\n",
        "  if not metrics_list:\n",
        "    return {}\n",
        "  rtn = {task:0. for task in metrics_list[0]}\n",
        "  for metric in metrics_list:\n",
        "    for task in metric:\n",
        "      rtn[task] += metric[task] / float(n)\n",
        "  return rtn \n",
        "\n",
        "def load_dataset(batch_size):\n",
        "  train = tfds.load('celeb_a', split='train')\n",
        "  resized_train = train.map(\n",
        "      lambda d: (d['attributes'], tf.image.resize(tf.image.convert_image_dtype(d['image'], tf.float32), [64, 64])))\n",
        "  final_train = resized_train.shuffle(\n",
        "      buffer_size=TRAIN_SIZE, seed=SEED,\n",
        "      reshuffle_each_iteration=True).batch(batch_size)\n",
        "\n",
        "  valid = tfds.load('celeb_a', split='validation')\n",
        "  resized_valid = valid.map(\n",
        "    lambda d: (d['attributes'], tf.image.resize(tf.image.convert_image_dtype(d['image'], tf.float32), [64, 64])))\n",
        "  final_valid = resized_valid.batch(batch_size)\n",
        "\n",
        "  test = tfds.load('celeb_a', split='test')\n",
        "  resized_test = test.map(lambda d: (d['attributes'], tf.image.resize(tf.image.convert_image_dtype(d['image'], tf.float32), [64, 64])))\n",
        "  final_test = resized_test.batch(batch_size)\n",
        "\n",
        "  Dataset = namedtuple('Dataset', ['train', 'valid', 'test'])\n",
        "  return Dataset(final_train, final_valid, final_test)\n",
        "\n",
        "def get_uncertainty_weights():\n",
        "  uncertainty_weights = {}\n",
        "  global shadow_uncertainty\n",
        "  if shadow_uncertainty is None:\n",
        "    shadow_uncertainty = tf.Variable(1.0)\n",
        "  uncertainty_weights['5_o_Clock_Shadow'] = shadow_uncertainty\n",
        "  global black_hair_uncertainty\n",
        "  if black_hair_uncertainty is None:\n",
        "    black_hair_uncertainty = tf.Variable(1.0)\n",
        "  uncertainty_weights['Black_Hair'] = black_hair_uncertainty\n",
        "  global blond_hair_uncertainty\n",
        "  if blond_hair_uncertainty is None:\n",
        "    blond_hair_uncertainty = tf.Variable(1.0)\n",
        "  uncertainty_weights['Blond_Hair'] = blond_hair_uncertainty\n",
        "  global brown_hair_uncertainty\n",
        "  if brown_hair_uncertainty is None:\n",
        "    brown_hair_uncertainty = tf.Variable(1.0)\n",
        "  uncertainty_weights['Brown_Hair'] = brown_hair_uncertainty\n",
        "  global goatee_uncertainty\n",
        "  if goatee_uncertainty is None:\n",
        "    goatee_uncertainty = tf.Variable(1.0)\n",
        "  uncertainty_weights['Goatee'] = goatee_uncertainty\n",
        "  global mustache_uncertainty\n",
        "  if mustache_uncertainty is None:\n",
        "    mustache_uncertainty = tf.Variable(1.0)\n",
        "  uncertainty_weights['Mustache'] = mustache_uncertainty\n",
        "  global no_beard_uncertainty\n",
        "  if no_beard_uncertainty is None:\n",
        "    no_beard_uncertainty = tf.Variable(1.0)\n",
        "  uncertainty_weights['No_Beard'] = no_beard_uncertainty\n",
        "  global rosy_cheeks_uncertainty\n",
        "  if rosy_cheeks_uncertainty is None:\n",
        "    rosy_cheeks_uncertainty = tf.Variable(1.0)\n",
        "  uncertainty_weights['Rosy_Cheeks'] = rosy_cheeks_uncertainty\n",
        "  global wearing_hat_uncertainty\n",
        "  if wearing_hat_uncertainty is None:\n",
        "    wearing_hat_uncertainty = tf.Variable(1.0)\n",
        "  uncertainty_weights['Wearing_Hat'] = wearing_hat_uncertainty\n",
        "  return uncertainty_weights\n",
        "\n",
        "def init_uncertainty_weights():\n",
        "  global shadow_uncertainty\n",
        "  shadow_uncertainty = None\n",
        "  global black_hair_uncertainty\n",
        "  black_hair_uncertainty = None\n",
        "  global blond_hair_uncertainty\n",
        "  blond_hair_uncertainty = None\n",
        "  global brown_hair_uncertainty\n",
        "  brown_hair_uncertainty = None\n",
        "  global goatee_uncertainty\n",
        "  goatee_uncertainty = None\n",
        "  global mustache_uncertainty\n",
        "  mustache_uncertainty = None\n",
        "  global no_beard_uncertainty\n",
        "  no_beard_uncertainty = None\n",
        "  global rosy_cheeks_uncertainty\n",
        "  rosy_cheeks_uncertainty = None\n",
        "  global wearing_hat_uncertainty\n",
        "  wearing_hat_uncertainty = None\n",
        "\n",
        "def init_gradnorm_weights():\n",
        "  global shadow_gradnorm\n",
        "  shadow_gradnorm = None\n",
        "  global black_hair_gradnorm\n",
        "  black_hair_gradnorm = None \n",
        "  global blond_hair_gradnorm\n",
        "  blond_hair_gradnorm = None\n",
        "  global brown_hair_gradnorm\n",
        "  brown_hair_gradnorm = None\n",
        "  global goatee_gradnorm\n",
        "  goatee_gradnorm = None\n",
        "  global mustache_gradnorm\n",
        "  mustache_gradnorm = None\n",
        "  global no_beard_gradnorm\n",
        "  no_beard_gradnorm = None\n",
        "  global rosy_cheeks_gradnorm\n",
        "  rosy_cheeks_gradnorm = None\n",
        "  global wearing_hat_gradnorm\n",
        "  wearing_hat_gradnorm = None\n",
        "\n",
        "def init_gradnorm_weights():\n",
        "  global shadow_gradnorm\n",
        "  shadow_gradnorm = None\n",
        "  global black_hair_gradnorm\n",
        "  black_hair_gradnorm = None \n",
        "  global blond_hair_gradnorm\n",
        "  blond_hair_gradnorm = None\n",
        "  global brown_hair_gradnorm\n",
        "  brown_hair_gradnorm = None\n",
        "  global goatee_gradnorm\n",
        "  goatee_gradnorm = None\n",
        "  global mustache_gradnorm\n",
        "  mustache_gradnorm = None\n",
        "  global no_beard_gradnorm\n",
        "  no_beard_gradnorm = None\n",
        "  global rosy_cheeks_gradnorm\n",
        "  rosy_cheeks_gradnorm = None\n",
        "  global wearing_hat_gradnorm\n",
        "  wearing_hat_gradnorm = None\n",
        "\n",
        "def fetch_gradnorm_weights():\n",
        "  gradnorm_weights = {}\n",
        "  global shadow_gradnorm\n",
        "  if shadow_gradnorm is None:\n",
        "    shadow_gradnorm = tf.Variable(1.0)\n",
        "  gradnorm_weights['5_o_Clock_Shadow'] = shadow_gradnorm\n",
        "  global black_hair_gradnorm\n",
        "  if black_hair_gradnorm is None:\n",
        "    black_hair_gradnorm = tf.Variable(1.0)\n",
        "  gradnorm_weights['Black_Hair'] = black_hair_gradnorm\n",
        "  global blond_hair_gradnorm\n",
        "  if blond_hair_gradnorm is None:\n",
        "    blond_hair_gradnorm = tf.Variable(1.0)\n",
        "  gradnorm_weights['Blond_Hair'] = blond_hair_gradnorm\n",
        "  global brown_hair_gradnorm\n",
        "  if brown_hair_gradnorm is None:\n",
        "    brown_hair_gradnorm = tf.Variable(1.0)\n",
        "  gradnorm_weights['Brown_Hair'] = brown_hair_gradnorm\n",
        "  global goatee_gradnorm\n",
        "  if goatee_gradnorm is None:\n",
        "    goatee_gradnorm = tf.Variable(1.0)\n",
        "  gradnorm_weights['Goatee'] = goatee_gradnorm\n",
        "  global mustache_gradnorm\n",
        "  if mustache_gradnorm is None:\n",
        "    mustache_gradnorm = tf.Variable(1.0)\n",
        "  gradnorm_weights['Mustache'] = mustache_gradnorm\n",
        "  global no_beard_gradnorm\n",
        "  if no_beard_gradnorm is None:\n",
        "    no_beard_gradnorm = tf.Variable(1.0)\n",
        "  gradnorm_weights['No_Beard'] = no_beard_gradnorm\n",
        "  global rosy_cheeks_gradnorm\n",
        "  if rosy_cheeks_gradnorm is None:\n",
        "    rosy_cheeks_gradnorm = tf.Variable(1.0)\n",
        "  gradnorm_weights['Rosy_Cheeks'] = rosy_cheeks_gradnorm\n",
        "  global wearing_hat_gradnorm\n",
        "  if wearing_hat_gradnorm is None:\n",
        "    wearing_hat_gradnorm = tf.Variable(1.0)\n",
        "  gradnorm_weights['Wearing_Hat'] = wearing_hat_gradnorm\n",
        "  return gradnorm_weights\n",
        "\n",
        "def init_gradnorm_l0():\n",
        "  global shadow_loss\n",
        "  shadow_loss = None\n",
        "  global black_hair_loss\n",
        "  black_hair_loss = None\n",
        "  global blond_hair_loss\n",
        "  blond_hair_loss = None\n",
        "  global brown_hair_loss\n",
        "  brown_hair_loss = None\n",
        "  global goatee_loss\n",
        "  goatee_loss = None\n",
        "  global mustache_loss\n",
        "  mustache_loss = None\n",
        "  global no_beard_loss\n",
        "  no_beard_loss = None\n",
        "  global rosy_cheeks_loss\n",
        "  rosy_cheeks_loss = None\n",
        "  global wearing_hat_loss\n",
        "  wearing_hat_loss = None\n",
        "\n",
        "def fetch_gradnorm_l0(losses):\n",
        "  gradnorm_l0 = {}\n",
        "  global shadow_loss\n",
        "  if shadow_loss is None:\n",
        "    if '5_o_Clock_Shadow' in losses:\n",
        "      loss_val = losses['5_o_Clock_Shadow']\n",
        "    else:\n",
        "      loss_val = 0.\n",
        "    shadow_loss = tf.Variable(loss_val)\n",
        "  gradnorm_l0['5_o_Clock_Shadow'] = shadow_loss\n",
        "  global black_hair_loss\n",
        "  if black_hair_loss is None:\n",
        "    if 'Black_Hair' in losses:\n",
        "      loss_val = losses['Black_Hair']\n",
        "    else:\n",
        "      loss_val = 0.\n",
        "    black_hair_loss = tf.Variable(loss_val)\n",
        "  gradnorm_l0['Black_Hair'] = black_hair_loss\n",
        "  global blond_hair_loss\n",
        "  if blond_hair_loss is None:\n",
        "    if 'Blond_Hair' in losses:\n",
        "      loss_val = losses['Blond_Hair']\n",
        "    else:\n",
        "      loss_val = 0.\n",
        "    blond_hair_loss = tf.Variable(loss_val)\n",
        "  gradnorm_l0['Blond_Hair'] = blond_hair_loss\n",
        "  global brown_hair_loss\n",
        "  if brown_hair_loss is None:\n",
        "    if 'Brown_Hair' in losses:\n",
        "      loss_val = losses['Brown_Hair']\n",
        "    else:\n",
        "      loss_val = 0.\n",
        "    brown_hair_loss = tf.Variable(loss_val)\n",
        "  gradnorm_l0['Brown_Hair'] = brown_hair_loss\n",
        "  global goatee_loss\n",
        "  if goatee_loss is None:\n",
        "    if 'Goatee' in losses:\n",
        "      loss_val = losses['Goatee']\n",
        "    else:\n",
        "      loss_val = 0.\n",
        "    goatee_loss = tf.Variable(loss_val)\n",
        "  gradnorm_l0['Goatee'] = goatee_loss\n",
        "  global mustache_loss\n",
        "  if mustache_loss is None:\n",
        "    if 'Mustache' in losses:\n",
        "      loss_val = losses['Mustache']\n",
        "    else:\n",
        "      loss_val = 0.\n",
        "    mustache_loss = tf.Variable(loss_val)\n",
        "  gradnorm_l0['Mustache'] = mustache_loss\n",
        "  global no_beard_loss\n",
        "  if no_beard_loss is None:\n",
        "    if 'No_Beard' in losses:\n",
        "      loss_val = losses['No_Beard']\n",
        "    else:\n",
        "      loss_val = 0.\n",
        "    no_beard_loss = tf.Variable(loss_val)\n",
        "  gradnorm_l0['No_Beard'] = no_beard_loss\n",
        "  global rosy_cheeks_loss\n",
        "  if rosy_cheeks_loss is None:\n",
        "    if 'Rosy_Cheeks' in losses:\n",
        "      loss_val = losses['Rosy_Cheeks']\n",
        "    else:\n",
        "      loss_val = 0.\n",
        "    rosy_cheeks_loss = tf.Variable(loss_val)\n",
        "  gradnorm_l0['Rosy_Cheeks'] = rosy_cheeks_loss\n",
        "  global wearing_hat_loss\n",
        "  if wearing_hat_loss is None:\n",
        "    if 'Wearing_Hat' in losses:\n",
        "      loss_val = losses['Wearing_Hat']\n",
        "    else:\n",
        "      loss_val = 0.\n",
        "    wearing_hat_loss = tf.Variable(loss_val)\n",
        "  gradnorm_l0['Wearing_Hat'] = wearing_hat_loss\n",
        "  return gradnorm_l0\n",
        "\n",
        "def compute_gradnorm_losses(losses, gradnorm_l0, gradnorms, expected_gradnorm):\n",
        "  task_li = {}\n",
        "  for task in FLAGS.tasks:\n",
        "    task_li[task] = losses[task] / gradnorm_l0[task]\n",
        "  li_expected = tf.reduce_mean(list(task_li.values()))\n",
        "\n",
        "  gradnorm_loss = {}\n",
        "  for task in FLAGS.tasks:\n",
        "    task_ri = tf.math.pow(task_li[task] / li_expected, params.alpha)\n",
        "    gradnorm_loss[task] = tf.norm(gradnorms[task] - tf.stop_gradient(expected_gradnorm*task_ri), ord=1)\n",
        "  total_gradnorm_loss = tf.reduce_sum(list(gradnorm_loss.values()))\n",
        "  return total_gradnorm_loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jSJ7yW8HX8IU"
      },
      "source": [
        "def train(params):\n",
        "  print(params)\n",
        "\n",
        "  ResBase = ResNet18()\n",
        "  ResTowers = {task:AttributeDecoder() for task in FLAGS.tasks}\n",
        "\n",
        "  dataset = load_dataset(FLAGS.batch_size)\n",
        "  global_step = tf.Variable(0, trainable=False)\n",
        "  optimizer = tf.keras.optimizers.SGD(params.lr, momentum=0.9)\n",
        "  if 'pcgrad' in FLAGS.method:\n",
        "    lr_var = tf.Variable(params.lr)\n",
        "    old_optimizer = tf.compat.v1.train.MomentumOptimizer(lr_var, momentum=0.9)\n",
        "    optimizer = PCGrad(tf.compat.v1.train.MomentumOptimizer(lr_var, momentum=0.9))\n",
        "\n",
        "  @tf.function()\n",
        "  def train_step(input, labels, first_step=False):\n",
        "    \"\"\"This is TAG.\"\"\"\n",
        "    with tf.GradientTape(persistent=True) as tape:\n",
        "      rep = ResBase(input, training=True)\n",
        "      preds = {task:model(rep, training=True) for (task, model) in ResTowers.items()}\n",
        "      losses = {task: tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(\n",
        "                      labels=labels[task], \n",
        "                      logits=preds[task]))\n",
        "                for task in labels}\n",
        "      loss = tf.add_n(list(losses.values()))\n",
        "\n",
        "      # Compute the gradient of the task-specific loss w.r.t. the shared base.\n",
        "      task_gains = {}\n",
        "      task_permutations = permute(losses)\n",
        "      combined_task_gradients = [(combined_task, tape.gradient(task_permutations[combined_task], ResBase.trainable_weights)) for combined_task in task_permutations]\n",
        "\n",
        "    for combined_task, task_gradient in combined_task_gradients:\n",
        "      if first_step:\n",
        "        base_update = [optimizer.lr*grad for grad in task_gradient]\n",
        "        base_updated = [param - update for param,update in zip(ResBase.trainable_weights, base_update)]\n",
        "      else:\n",
        "        base_update = [(optimizer._momentum*optimizer.get_slot(param, 'momentum') - optimizer.lr*grad) \n",
        "                        for param, grad in zip(ResBase.trainable_weights, task_gradient)]\n",
        "        base_updated = [param + update for param, update in zip(ResBase.trainable_weights, base_update)]\n",
        "      task_update_rep = base_step(input, base_updated)\n",
        "      task_update_preds = {task:model(task_update_rep, training=True) for (task, model) in ResTowers.items()}\n",
        "      task_update_losses = {task: tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(\n",
        "                                  labels=labels[task], \n",
        "                                  logits=task_update_preds[task]))\n",
        "                            for task in labels}\n",
        "      task_gain = {task:(1.0 - task_update_losses[task]/losses[task])/optimizer.lr for task in FLAGS.tasks}\n",
        "      task_gains[combined_task] = task_gain\n",
        "\n",
        "    # DO NOT apply Nesterov in normal mtl training.\n",
        "    for task,model in ResTowers.items():\n",
        "      task_grads = tape.gradient(loss, model.trainable_weights)\n",
        "      optimizer.apply_gradients(zip(task_grads, model.trainable_weights))\n",
        "    \n",
        "    # Apply the traditional MTL update since this is a normal train step.\n",
        "    base_grads = tape.gradient(loss, ResBase.trainable_weights)\n",
        "    optimizer.apply_gradients(zip(base_grads, ResBase.trainable_weights))\n",
        "\n",
        "    global_step.assign_add(1)\n",
        "    return losses, task_gains\n",
        "\n",
        "  @tf.function()\n",
        "  def cosine_sim_train_step(input, labels, first_step=False):\n",
        "    with tf.GradientTape(persistent=True) as tape:\n",
        "      rep = ResBase(input, training=True)\n",
        "      preds = {task:model(rep, training=True) for (task, model) in ResTowers.items()}\n",
        "      losses = {task: tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(\n",
        "                      labels=labels[task], \n",
        "                      logits=preds[task]))\n",
        "                for task in labels}\n",
        "      loss = tf.add_n(list(losses.values()))\n",
        "\n",
        "      # Compute the gradient of the task-specific loss w.r.t. the shared base.\n",
        "      task_gains = {}\n",
        "      task_permutations = permute(losses)\n",
        "\n",
        "      task_gradients = {}\n",
        "      for task in FLAGS.tasks:\n",
        "        task_grads = tape.gradient(task_permutations[task], ResBase.trainable_weights)\n",
        "        task_grads = [grad for grad in task_grads if len(grad.shape) > 1]\n",
        "        task_gradients[task] = task_grads\n",
        "      combined_task_gradients = [(combined_task, tape.gradient(task_permutations[combined_task], ResBase.trainable_weights)) for combined_task in task_permutations]\n",
        "\n",
        "    for combined_task, task_gradient in combined_task_gradients:\n",
        "      task_gain = {}\n",
        "      for task in FLAGS.tasks:\n",
        "        filtered_grads = [grad for grad in task_gradient if len(grad.shape) > 1]\n",
        "        stacked_filtered = tf.concat([tf.reshape(grad, shape=[-1]) for grad in filtered_grads], axis=0)\n",
        "        stacked_task = tf.concat([tf.reshape(grad, shape=[-1]) for grad in task_gradients[task]], axis=0)\n",
        "        cosine_sim = tf.reduce_sum(tf.multiply(stacked_filtered, stacked_task)) / (tf.norm(stacked_filtered, ord=2) * tf.norm(stacked_task, ord=2))\n",
        "        task_gain[task] = cosine_sim\n",
        "      task_gains[combined_task] = task_gain\n",
        "\n",
        "    # DO NOT apply Nesterov in normal mtl training.\n",
        "    for task,model in ResTowers.items():\n",
        "      task_grads = tape.gradient(loss, model.trainable_weights)\n",
        "      optimizer.apply_gradients(zip(task_grads, model.trainable_weights))\n",
        "    \n",
        "    # Apply the traditional MTL update since this is a normal train step.\n",
        "    base_grads = tape.gradient(loss, ResBase.trainable_weights)\n",
        "    optimizer.apply_gradients(zip(base_grads, ResBase.trainable_weights))\n",
        "\n",
        "    global_step.assign_add(1)\n",
        "    return losses, task_gains\n",
        "\n",
        "  @tf.function()\n",
        "  def train_fast_step(input, labels, first_step=False):\n",
        "    \"\"\"Call this function to evaluate task groupings. It's faster.\"\"\"\n",
        "    with tf.GradientTape(persistent=True) as tape:\n",
        "      rep = ResBase(input, training=True)\n",
        "      preds = {task:model(rep, training=True) for (task, model) in ResTowers.items()}\n",
        "      losses = {task: tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(\n",
        "                      labels=labels[task], \n",
        "                      logits=preds[task]))\n",
        "                for task in labels}\n",
        "      loss = tf.add_n(list(losses.values()))\n",
        "\n",
        "    # DO NOT apply Nesterov in normal mtl training.\n",
        "    for task,model in ResTowers.items():\n",
        "      task_grads = tape.gradient(loss, model.trainable_weights)\n",
        "      optimizer.apply_gradients(zip(task_grads, model.trainable_weights))\n",
        "    \n",
        "    # Apply the traditional MTL update since this is a normal train step.\n",
        "    base_grads = tape.gradient(loss, ResBase.trainable_weights)\n",
        "    optimizer.apply_gradients(zip(base_grads, ResBase.trainable_weights))\n",
        "\n",
        "    global_step.assign_add(1)\n",
        "    return losses, {}\n",
        "\n",
        "  @tf.function()\n",
        "  def train_uncertainty_step(input, labels, first_step=False):\n",
        "    with tf.GradientTape(persistent=True) as tape:\n",
        "      rep = ResBase(input, training=True)\n",
        "      preds = {task:model(rep, training=True) for (task, model) in ResTowers.items()}\n",
        "      losses = {task: tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(\n",
        "                      labels=labels[task], \n",
        "                      logits=preds[task]))\n",
        "                for task in labels}\n",
        "\n",
        "      uncertainty_weights = get_uncertainty_weights()\n",
        "      for task in FLAGS.tasks:\n",
        "        clip_uncertainty = tf.clip_by_value(uncertainty_weights[task], 0.01, 10.0)\n",
        "        losses[task] = losses[task] / tf.exp(2 * clip_uncertainty) + clip_uncertainty\n",
        "      loss = tf.add_n(list(losses.values()))\n",
        "\n",
        "    # DO NOT apply Nesterov in normal mtl training.\n",
        "    for task,model in ResTowers.items():\n",
        "      task_grads = tape.gradient(loss, model.trainable_weights)\n",
        "      optimizer.apply_gradients(zip(task_grads, model.trainable_weights))\n",
        "    \n",
        "    # Apply the traditional MTL update since this is a normal train step.\n",
        "    base_grads = tape.gradient(loss, ResBase.trainable_weights)\n",
        "    optimizer.apply_gradients(zip(base_grads, ResBase.trainable_weights))\n",
        "\n",
        "    # Update the uncertainty weight variables.\n",
        "    uncertainty_gradients = [val for val in tape.gradient(loss, list(uncertainty_weights.values()))]\n",
        "    optimizer.apply_gradients(zip(uncertainty_gradients, list(uncertainty_weights.values())))\n",
        "\n",
        "    global_step.assign_add(1)\n",
        "    return losses, {}\n",
        "\n",
        "  @tf.function()\n",
        "  def train_gradnorm_step(input, labels, first_step=False):\n",
        "    with tf.GradientTape(persistent=True) as tape:\n",
        "      rep = ResBase(input, training=True)\n",
        "      preds = {task:model(rep, training=True) for (task, model) in ResTowers.items()}\n",
        "      losses = {task: tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(\n",
        "                      labels=labels[task], \n",
        "                      logits=preds[task]))\n",
        "                for task in labels}\n",
        "\n",
        "      # Essentially gradnorm weights.\n",
        "      gradnorm_weights = fetch_gradnorm_weights()\n",
        "      post_weight_losses = {task:gradnorm_weights[task]*losses[task] for task in losses}\n",
        "      gradnorms = {}\n",
        "      for task in FLAGS.tasks:\n",
        "        gradnorms[task] = tf.norm(tape.gradient(post_weight_losses[task], ResBase.trainable_weights[-3]), ord=2)\n",
        "      expected_gradnorm = tf.add_n(list(gradnorms.values())) / len(FLAGS.tasks)\n",
        "      loss = tf.add_n(list(post_weight_losses.values()))\n",
        "      \n",
        "      gradnorm_l0 = fetch_gradnorm_l0(losses)\n",
        "      gradnorm_losses = compute_gradnorm_losses(losses, gradnorm_l0, gradnorms, expected_gradnorm)\n",
        "\n",
        "    # Update gradnorm weights.\n",
        "    gradnorm_weight_grads = tape.gradient(gradnorm_losses, gradnorm_weights.values())\n",
        "    optimizer.apply_gradients(zip(gradnorm_weight_grads, list(gradnorm_weights.values())))\n",
        "\n",
        "    # DO NOT apply Nesterov in normal mtl training.\n",
        "    for task,model in ResTowers.items():\n",
        "      task_grads = tape.gradient(loss, model.trainable_weights)\n",
        "      optimizer.apply_gradients(zip(task_grads, model.trainable_weights))\n",
        "    \n",
        "    # Apply the traditional MTL update since this is a normal train step.\n",
        "    base_grads = tape.gradient(loss, ResBase.trainable_weights)\n",
        "    optimizer.apply_gradients(zip(base_grads, ResBase.trainable_weights))\n",
        "\n",
        "    # Clip between 0.1 and 10\n",
        "    for task in FLAGS.tasks:\n",
        "      gradnorm_weights[task].assign(tf.clip_by_value(gradnorm_weights[task], 0.1, 10.0))\n",
        "\n",
        "    # Renormalize GradNorm Weights\n",
        "    scale = len(FLAGS.tasks) / tf.reduce_sum(list(gradnorm_weights.values()))\n",
        "    for task in FLAGS.tasks:\n",
        "      gradnorm_weights[task].assign(scale*gradnorm_weights[task])\n",
        "\n",
        "    global_step.assign_add(1)\n",
        "    return losses, {}\n",
        "\n",
        "  @tf.function()\n",
        "  def train_pcgrad_step(input, labels, first_step=False):\n",
        "    rep = ResBase(input, training=True)\n",
        "    preds = {task:model(rep, training=True) for (task, model) in ResTowers.items()}\n",
        "    losses = {task: tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(\n",
        "                    labels=labels[task], \n",
        "                    logits=preds[task]))\n",
        "              for task in labels}\n",
        "    uncertainty_weights = get_uncertainty_weights()\n",
        "    for task in FLAGS.tasks:\n",
        "      clip_uncertainty = tf.clip_by_value(uncertainty_weights[task], 0.01, 10.0)\n",
        "      losses[task] = losses[task] / tf.exp(2 * clip_uncertainty) + clip_uncertainty\n",
        "    loss = tf.add_n(list(losses.values()))\n",
        "\n",
        "    base_gradvars = optimizer.compute_gradients(list(losses.values()), ResBase.trainable_weights)\n",
        "    task_gradvars = [optimizer.compute_gradients([losses[task]], model.trainable_weights) for (task,model) in ResTowers.items()]\n",
        "\n",
        "    old_optimizer.apply_gradients(base_gradvars)\n",
        "    for gv in task_gradvars:\n",
        "      old_optimizer.apply_gradients(gv)\n",
        "\n",
        "    # Update the uncertainty weight variables.\n",
        "    uw_gv = old_optimizer.compute_gradients(loss, list(uncertainty_weights.values()))\n",
        "    old_optimizer.apply_gradients(uw_gv)\n",
        "\n",
        "    global_step.assign_add(1)\n",
        "    return losses, {}\n",
        "\n",
        "  @tf.function()\n",
        "  def eval_step(input, labels):\n",
        "    rep = ResBase(input)\n",
        "    preds = {task:ResTowers[task](rep) for (task, model) in ResTowers.items()}\n",
        "    int_preds = {task:tf.math.argmax(preds[task], 1, tf.dtypes.int32) for task in labels}\n",
        "    int_labels = {task:tf.math.argmax(labels[task], 1, tf.dtypes.int32) for task in labels}\n",
        "    losses = {task: tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(\n",
        "                labels=tf.cast(labels[task], tf.float32), \n",
        "                logits=preds[task])) \n",
        "              for task in labels}\n",
        "    accuracies = {task:tf.math.count_nonzero(tf.equal(int_preds[task], int_labels[task])) for task in labels}\n",
        "    Eval = namedtuple('Eval', ['losses', 'accuracies'])\n",
        "    return Eval(losses, accuracies)\n",
        "\n",
        "  # Training Loop.\n",
        "  metrics = {'train_loss': [], 'eval_loss': [], 'eval_acc': []}\n",
        "  gradient_metrics = {task:[] for task in permute_list(FLAGS.tasks)}\n",
        "  final_metrics = {'train_loss': [], 'eval_loss': [], 'eval_acc': []}\n",
        "  model_params = []\n",
        "\n",
        "  end = None\n",
        "  for step in range(FLAGS.steps):\n",
        "    if end:\n",
        "      print(f'Differnece in time: {end - start}')\n",
        "    start = time.time()\n",
        "    print('epoch: {}'.format(step))\n",
        "    if \"pcgrad\" not in FLAGS.method:\n",
        "      decay_lr(step, optimizer) # Halve the learning rate every 30 steps.\n",
        "    else:\n",
        "      decay_pcgrad_lr(step, lr_var)\n",
        "    batch_train_loss = {task:0. for task in FLAGS.tasks}\n",
        "    batch_grad_metrics = {combined_task:{task:0. for task in FLAGS.tasks} for combined_task in gradient_metrics}\n",
        "    for labels, img in dataset.train:\n",
        "      labels = {task:tf.keras.utils.to_categorical(labels[task], num_classes=2) for task in labels if task in FLAGS.tasks}\n",
        "      if FLAGS.method == 'mtl': # Full TAG.\n",
        "         losses, task_gains = train_step(img, labels, first_step=(len(optimizer.variables()) == 0))\n",
        "      elif FLAGS.method == 'cosine_sim_mtl': # Cosine similarity computation.\n",
        "         losses, task_gains = cosine_sim_train_step(img, labels, first_step=(len(optimizer.variables()) == 0))\n",
        "      elif FLAGS.method == 'fast_mtl': # Runs w/o TAG computation.\n",
        "         losses, task_gains = train_fast_step(img, labels, first_step=(len(optimizer.variables()) == 0))\n",
        "      elif FLAGS.method == 'uncertainty_mtl': # Runs w/ uncertainty weights.\n",
        "         losses, task_gains = train_uncertainty_step(img, labels, first_step=(len(optimizer.variables()) == 0))\n",
        "      elif FLAGS.method == 'gradnorm_mtl': # Runs w/ gradnorm weights.\n",
        "         losses, task_gains = train_gradnorm_step(img, labels, first_step=(len(optimizer.variables()) == 0))\n",
        "      elif FLAGS.method == 'pcgrad_mtl': # Runs w/ PCGrad Optimizer and UW.\n",
        "        losses, task_gains = train_pcgrad_step(img, labels, first_step=(len(optimizer.variables()) == 0))\n",
        "      else:\n",
        "        raise Exception(\"Unrecognized Method Selected.\")\n",
        "      \n",
        "      # Record batch-level training and gradient metrics. \n",
        "      for combined_task,task_gain_map in task_gains.items():\n",
        "        for task,gain in task_gain_map.items():\n",
        "          batch_grad_metrics[combined_task][task] += gain.numpy() / (math.ceil(TRAIN_SIZE / FLAGS.batch_size))\n",
        "      for task,loss in losses.items():\n",
        "        batch_train_loss[task] += loss.numpy() / (math.ceil(TRAIN_SIZE / FLAGS.batch_size))\n",
        "\n",
        "    # Record epoch-level training and gradient metrics.\n",
        "    add_average(metrics['train_loss'], batch_train_loss, METRICS_AVERAGE)\n",
        "    for combined_task,task_gain_map in batch_grad_metrics.items():\n",
        "      gradient_metrics[combined_task].append(task_gain_map)\n",
        "\n",
        "    batch_eval_loss = {task:0. for task in FLAGS.tasks}\n",
        "    batch_eval_acc = {task:0. for task in FLAGS.tasks}\n",
        "    for labels, img in dataset.test if FLAGS.eval == 'test' else dataset.valid:\n",
        "      labels = {task:tf.keras.utils.to_categorical(labels[task], num_classes=2) for task in labels if task in FLAGS.tasks}\n",
        "      eval_metrics = eval_step(img, labels)\n",
        "      for task in FLAGS.tasks:\n",
        "        EVAL_SIZE = TEST_SIZE if FLAGS.eval == 'test' else VALID_SIZE\n",
        "        batch_eval_loss[task] += eval_metrics.losses[task].numpy() / (math.ceil(EVAL_SIZE / FLAGS.batch_size))\n",
        "        batch_eval_acc[task] += eval_metrics.accuracies[task].numpy() / EVAL_SIZE\n",
        "    add_average(metrics['eval_loss'], batch_eval_loss, METRICS_AVERAGE)\n",
        "    add_average(metrics['eval_acc'], batch_eval_acc, METRICS_AVERAGE)\n",
        "\n",
        "    for metric in metrics:\n",
        "      final_metrics[metric].append(compute_average(metrics[metric], METRICS_AVERAGE))\n",
        "\n",
        "    # Save past EARLY_STOP sets of parameters.\n",
        "    cur_params = [('base', copy.deepcopy(ResBase.trainable_weights), copy.deepcopy(ResBase.non_trainable_weights))] + \\\n",
        "        [(task, copy.deepcopy(tower.trainable_weights), copy.deepcopy(tower.non_trainable_weights)) for task,tower in ResTowers.items()]\n",
        "    model_params.append(tuple(cur_params))\n",
        "    \n",
        "    # Early stopping. If Validation loss hasn't increased for the past 10 epochs..\n",
        "    EARLY_STOP = 11\n",
        "    if step > EARLY_STOP - 1 and all([sum(final_metrics['eval_loss'][-EARLY_STOP].values()) < sum(final_metrics['eval_loss'][-i].values()) for i in range(1, EARLY_STOP)]):\n",
        "      print('Validation loss has not improved for past 10 epochs. Stopping at epoch {}'.format(step))\n",
        "\n",
        "      # Reload best weights..\n",
        "      for task, trainables, non_trainables in model_params[-EARLY_STOP]:\n",
        "        if task == 'base':\n",
        "          for param, trainable in zip(ResBase.trainable_weights, trainables):\n",
        "            param.assign(trainable)\n",
        "          for param, non_trainable in zip(ResBase.non_trainable_weights, non_trainables):\n",
        "            param.assign(non_trainable)\n",
        "        else:\n",
        "          for param, trainable in zip(ResTowers[task].trainable_weights, trainables):\n",
        "            param.assign(trainable)\n",
        "          for param, non_trainable in zip(ResTowers[task].non_trainable_weights, non_trainables):\n",
        "            param.assign(non_trainable)\n",
        "\n",
        "      # Evaluate on the test set.\n",
        "      batch_test_acc = {task:0. for task in FLAGS.tasks}\n",
        "      batch_test_loss = {task:0. for task in FLAGS.tasks}\n",
        "      for labels, img in dataset.test:\n",
        "        labels = {task:tf.keras.utils.to_categorical(labels[task], num_classes=2) for task in labels if task in FLAGS.tasks}\n",
        "        test_metrics = eval_step(img, labels)\n",
        "        for task in FLAGS.tasks:\n",
        "          EVAL_SIZE = TEST_SIZE if FLAGS.eval == 'test' else VALID_SIZE\n",
        "          batch_test_loss[task] += test_metrics.losses[task].numpy() / (math.ceil(EVAL_SIZE / FLAGS.batch_size))\n",
        "          batch_test_acc[task] += test_metrics.accuracies[task].numpy() / EVAL_SIZE\n",
        "\n",
        "      print_test_acc = \"\\n\".join([\"{}: {:.2f}\".format(task, 100.0*metric) for task, metric in batch_test_acc.items()])\n",
        "      print_test_loss = \"\\n\".join([\"{}: {:.4f}\".format(task, metric) for task, metric in batch_test_loss.items()])\n",
        "      print(\"Test Accuracy:\\n{}\\n\".format(print_test_acc))\n",
        "      print(\"Test Loss:\\n{}\\n\".format(print_test_loss))\n",
        "\n",
        "      for task in gradient_metrics:\n",
        "        gradient_metrics[task] = gradient_metrics[task][:-1*(EARLY_STOP - 1)]\n",
        "\n",
        "      return final_metrics, gradient_metrics\n",
        "\n",
        "    elif len(model_params) == EARLY_STOP:\n",
        "      model_params.pop(0)\n",
        "\n",
        "    print_train_loss = \"\\n\".join([\"{}: {:.4f}\".format(task, metric) for task, metric in final_metrics['train_loss'][-1].items()])\n",
        "    print(\"Train Loss:\\n{}\\n\".format(print_train_loss))\n",
        "\n",
        "    print(\"grad metrics for fun: {}\".format(gradient_metrics))\n",
        "\n",
        "    print_eval_loss = \"\\n\".join([\"{}: {:.4f}\".format(task, metric) for task, metric in final_metrics['eval_loss'][-1].items()])\n",
        "    print(\"Eval Loss:\\n{}\\n\".format(print_eval_loss))\n",
        "    print_eval_acc = \"\\n\".join([\"{}: {:.2f}\".format(task, 100.0*metric) for task, metric in final_metrics['eval_acc'][-1].items()])\n",
        "    print(\"Eval Accuracy:\\n{}\\n\".format(print_eval_acc))\n",
        "    print(\"\\n-------------\\n\")\n",
        "    end = time.time()\n",
        "\n",
        "  return final_metrics, gradient_metrics"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cImVvSUgcGhD"
      },
      "source": [
        "Params = namedtuple(\"Params\", ['lr', 'alpha']) # Params can possibly be tuned, FLAGS can't be tuned.\n",
        "params = Params(lr=0.0005, alpha=0.1)\n",
        "FLAGS.steps = 100 # MOO: train for 100 epochs.\n",
        "FLAGS.batch_size = 256 # MOO: train with batch size = 256 \n",
        "FLAGS.eval = 'valid'\n",
        "FLAGS.method = 'fast_mtl'\n",
        "FLAGS.order = 1\n",
        "FLAGS.tasks = ['5_o_Clock_Shadow', 'Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Goatee', 'Mustache', 'No_Beard', 'Rosy_Cheeks', 'Wearing_Hat'] # 9 out of 40 attributes."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "urok0MhIvUER"
      },
      "source": [
        "if FLAGS.method == 'uncertainty_mtl' or 'pcgrad' in FLAGS.method:\n",
        "  init_uncertainty_weights()\n",
        "\n",
        "if FLAGS.method == 'gradnorm_mtl':\n",
        "  init_gradnorm_weights()\n",
        "  init_gradnorm_l0()\n",
        "\n",
        "# %%capture\n",
        "# run the model 1 time\n",
        "tf.compat.v1.reset_default_graph()\n",
        "eval_metrics, gradient_metrics = train(params)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}